#include <lobi/zipfile.h>

#include <assert.h>
#include <stdlib.h>
#include <string.h>

#include <zlib.h>

typedef unsigned char uchar_t; /* 8 bits unsigned (1 byte) */
typedef unsigned short ushort_t; /* 16 bits unsigned (2 bytes)*/
typedef unsigned int uint_t; /* 32 bits unsigned (4 bytes) */

/*
 * .zip file data structures
 */


#ifdef _MSC_VER
#pragma pack(push,1)
#define ATTR_PACKED
#elif defined(__GNUC__)
#define ATTR_PACKED __attribute__((packed))
#else
#error unknown compiler.
#endif

/***
ZIP format
Byte order: Little-endian

Overall zipfile format:
[Local file header + Compressed data [+ Extended local header]?]*
[Central directory]*
[End of central directory record]

compression method: (2 bytes)
          0 - The file is stored (no compression)
          1 - The file is Shrunk
          2 - The file is Reduced with compression factor 1
          3 - The file is Reduced with compression factor 2
          4 - The file is Reduced with compression factor 3
          5 - The file is Reduced with compression factor 4
          6 - The file is Imploded
          7 - Reserved for Tokenizing compression algorithm
          8 - The file is Deflated
***/

/*************************
Local file header:
Offset   Length   Contents
  0      4 bytes  Local file header signature (0x04034b50)
  4      2 bytes  Version needed to extract
  6      2 bytes  General purpose bit flag
  8      2 bytes  Compression method
 10      2 bytes  Last mod file time
 12      2 bytes  Last mod file date
 14      4 bytes  CRC-32
 18      4 bytes  Compressed size (n)
 22      4 bytes  Uncompressed size
 26      2 bytes  Filename length (f)
 28      2 bytes  Extra field length (e)
        (f)bytes  Filename
        (e)bytes  Extra field
        (n)bytes  Compressed data
*************************/
typedef struct local_file_header_t
{
	uint_t   sig;
	ushort_t extract_ver;
	ushort_t bit_flag;
	ushort_t compress_method;
	ushort_t file_time;
	ushort_t file_date;
	uint_t   crc;
	uint_t   compr_size;
	uint_t   uncompr_size;
	ushort_t filename_len;
	ushort_t extra_field_len;
} ATTR_PACKED local_file_header_t;

/**************************
Extended local header:
Offset   Length   Contents
  0      4 bytes  Extended Local file header signature (0x08074b50)
  4      4 bytes  CRC-32
  8      4 bytes  Compressed size
 12      4 bytes  Uncompressed size
**************************/
typedef struct extended_local_header_t
{
	uint_t  sig;
	uint_t  crc;
	uint_t  compr_size;
	uint_t  uncompr_size;
} ATTR_PACKED extended_local_header_t;

/**************************
Central directory:
Offset   Length   Contents
  0      4 bytes  Central file header signature (0x02014b50)
  4      2 bytes  Version made by
  6      2 bytes  Version needed to extract
  8      2 bytes  General purpose bit flag
 10      2 bytes  Compression method
 12      2 bytes  Last mod file time
 14      2 bytes  Last mod file date
 16      4 bytes  CRC-32
 20      4 bytes  Compressed size
 24      4 bytes  Uncompressed size
 28      2 bytes  Filename length (f)
 30      2 bytes  Extra field length (e)
 32      2 bytes  File comment length (c)
 34      2 bytes  Disk number start
 36      2 bytes  Internal file attributes
 38      4 bytes  External file attributes
 42      4 bytes  Relative offset of local header
 46     (f)bytes  Filename
        (e)bytes  Extra field
        (c)bytes  File comment
**************************/
typedef struct central_dirent_t
{
	uint_t   sig;
	ushort_t version_made;
	
	ushort_t extract_ver;
	ushort_t bit_flag;
	ushort_t compress_method;
	ushort_t file_time;
	ushort_t file_date;
	uint_t   crc;
	uint_t   compr_size;
	uint_t   uncompr_size;
	ushort_t filename_len;
	ushort_t extra_field_len;

	ushort_t file_comment_len;
	ushort_t disk_num;
	ushort_t internal_file_attr;
	uint_t   ext_file_attr;
	uint_t   rel_offs_local_hdr;

} ATTR_PACKED central_dirent_t;

/**************************
End of central directory record:
Offset   Length   Contents
  0      4 bytes  End of central dir signature (0x06054b50)
  4      2 bytes  Number of this disk
  6      2 bytes  Number of the disk with the start of the central directory
  8      2 bytes  Total number of entries in the central dir on this disk
 10      2 bytes  Total number of entries in the central dir
 12      4 bytes  Size of the central directory
 16      4 bytes  Offset of start of central directory with respect to 
				  the starting disk number
 20      2 bytes  zipfile comment length (c)
 22     (c)bytes  zipfile comment
**************************/
typedef struct central_dir_end_t
{
	uint_t   sig;
	ushort_t diskno;
	ushort_t diskno_with_central_dir;
	ushort_t num_entries_this_disk;
	ushort_t num_entries_central_dir;
	uint_t   central_dir_size;
	uint_t   central_dir_offset;
	ushort_t zipfile_comment_len;
} ATTR_PACKED central_dir_end_t;

#ifdef _MSC_VER
#pragma pack(pop)
#endif


/*
 * public functions
 */


zipfile_t* zipfile_init(const char* filename)
{
	zipfile_t* zipfile = malloc(sizeof(zipfile_t));
	central_dir_end_t central_dir_end;
	int cde_comment_len, i;

	/* assert size of built-in types for i386 */
	assert(sizeof(uchar_t) == 1 && sizeof(ushort_t) == 2 && sizeof(uint_t) == 4);

	/* assertion check for (packed) structure sizes here */
	assert(sizeof(local_file_header_t) == 30);
	assert(sizeof(extended_local_header_t) == 16);
	assert(sizeof(central_dirent_t) == 46);
	assert(sizeof(central_dir_end_t) == 22);

	/* assertion check for little-endian here */
	{
		ushort_t i = 0x4321;
		uchar_t* s = (char*) &i;
		assert(*s == 0x21);
	}
	
	zipfile->fp = fopen(filename,"rb");
	if (!zipfile->fp) return 0;

	/* seek central_dir_end struct, assume comment no longer than 256 */
	for (cde_comment_len = 0; cde_comment_len < 256; cde_comment_len++) {
		if (fseek(zipfile->fp,-((int)sizeof(central_dir_end_t)+cde_comment_len),SEEK_END) != 0) goto die;
		if (fread(&central_dir_end,sizeof(central_dir_end_t), 1,zipfile->fp) != 1) goto die;
		if (central_dir_end.sig == 0x06054b50) 
			break;
	}
	if (cde_comment_len == 256) { 
		printf("cant find zipfile end. comment longer than 256 chars\n");
		goto die;
	}

	zipfile->nfiles = central_dir_end.num_entries_central_dir;
	zipfile->entries = calloc(zipfile->nfiles, sizeof(zipfile_entry_t));

	if (fseek(zipfile->fp,-((int)sizeof(central_dir_end_t)+(int)cde_comment_len+(int)central_dir_end.central_dir_size),SEEK_END) != 0) goto die;

	for (i = 0; i < zipfile->nfiles; i++)
	{
		central_dirent_t central_dirent;
		char filename[1024];
		unsigned int len;

		/* seek to begin of central dir */
		if (fread(&central_dirent,sizeof(central_dirent_t),1,zipfile->fp) != 1) goto die;
		if (central_dirent.sig != 0x02014b50) goto die;
		
		if (fread(filename, 1, len = central_dirent.filename_len,zipfile->fp) != len) goto die;
		filename[len] = 0;

		zipfile->entries[i].filename = strdup(filename);
		zipfile->entries[i].offset = central_dirent.rel_offs_local_hdr;
		/* skip extra field */
		if (central_dirent.extra_field_len)
			fseek(zipfile->fp,central_dirent.extra_field_len, SEEK_CUR);

		/* check local header */
		{
			local_file_header_t local_file_header;
			long pos = ftell(zipfile->fp);

			fseek(zipfile->fp, zipfile->entries[i].offset, SEEK_SET);
			fread(&local_file_header,sizeof(local_file_header_t),1,zipfile->fp);
			if (local_file_header.sig != 0x04034b50) goto die;

			/* restore file pointer */
			fseek(zipfile->fp,pos,SEEK_SET);
		}
	}

	return zipfile;
die:
	perror(filename);
	zipfile_done(zipfile);
	return 0;
}


void zipfile_done(zipfile_t* zipfile)
{
	if (zipfile)
	{
		int i;

		if (zipfile->fp) { fclose(zipfile->fp); zipfile->fp = 0; }

		/* free filenames */
		for (i = 0; i < zipfile->nfiles; i++)
			if (zipfile->entries[i].filename) free(zipfile->entries[i].filename);
		free(zipfile->entries);

		free(zipfile);
	}
}


void* zipfile_slorp(zipfile_t* zipfile, const char* filename, unsigned int* length)
{
	if (zipfile && filename && length)
	{
		int i;

		for (i = 0; i < zipfile->nfiles; i++)
		{
			if (strcmp(zipfile->entries[i].filename,filename) == 0)
			{
				local_file_header_t local_file_header;

				fseek(zipfile->fp,zipfile->entries[i].offset,SEEK_SET);
				fread(&local_file_header,sizeof(local_file_header_t),1,zipfile->fp);
				if (local_file_header.sig != 0x04034b50) return 0;

				/* skip name */
				fseek(zipfile->fp,
					local_file_header.filename_len+local_file_header.extra_field_len,
					SEEK_CUR);

				switch(local_file_header.compress_method)
				{
					case 0: /* stored */
					{
						uint_t src_len = local_file_header.compr_size;
						uint_t dst_len = local_file_header.uncompr_size;
						uchar_t* src = malloc(src_len);

						assert(src_len == dst_len);
						fread(src,1,src_len,zipfile->fp);
						*length = dst_len;
						return src;
					}
					case 8: /* deflated */
					{
						uint_t src_len = local_file_header.compr_size;
						uint_t dst_len = local_file_header.uncompr_size;
						uchar_t* src = malloc(src_len);
						uchar_t* dst = malloc(dst_len);

						fread(src,1,src_len,zipfile->fp);
						*length = dst_len;

						/* inflate buffer */
						{
							int err;
							z_stream d_stream;

						    d_stream.zalloc = 0;
						    d_stream.zfree = 0;
						    d_stream.opaque = 0;

						    d_stream.next_in  = src;
							d_stream.avail_in = 0;
							d_stream.next_out = dst;
								
							/* skip .gz header, there is none in .zip files */
							err = inflateInit2(&d_stream,-MAX_WBITS);
							if (err != Z_OK) return 0;

	/* we use <= instead of < because there is no zip header */
    while (d_stream.total_out <= dst_len && d_stream.total_in <= src_len) {
        d_stream.avail_in = d_stream.avail_out = 1; /* force small buffers */
        err = inflate(&d_stream, Z_NO_FLUSH);
        if (err == Z_STREAM_END) break;
        if (err != Z_OK) return 0;
    }
							
							err = inflateEnd(&d_stream);
							if (err != Z_OK) return 0;
						}
						
						free(src);	
						return dst;
					}
					default:
						printf("zipfile_slorp():unsupported compression type\n");
						return 0;
				}
			}
		}

		return 0;
	}
	else
		return 0;
}


void* file_slorp(const char* filename, unsigned int* length)
{
	FILE* fp;
	unsigned char* data;

	if ((fp=fopen(filename,"rb")) == 0) return 0;
	fseek(fp,0,SEEK_END);
	*length = ftell(fp);
	data = malloc(*length);
	rewind(fp);
	fread(data,1,*length,fp);
	fclose(fp);

	return data;
}

